#include <Heap.h>
#include <Globals.h>
#include <MemoryManagement.h>
#include <MonitorDriver.h>
#include <ProcessManager.h>
#include <Process.h>

bool lessThan(Header *a, Header *b)
{
  return a->size < b->size;
}

Heap::Heap(u32int start, u32int end, u32int max, bool isKernel) :
  startAddress(start), endAddress(end), maxAddress(max), isKernel(isKernel)
{

  ASSERT(startAddress%PAGE_SIZE == 0);
  ASSERT(endAddress%PAGE_SIZE == 0);

  // Initialise the index.
  index = (OrderedArray<Header*,HEAP_INDEX_SIZE> *)startAddress;

  index->setPredicate(lessThan);

  // Shift the start address to resemble where we can start putting data.
  startAddress += sizeof(*index);

  // make sure startAddress is page-aligned.
  if (startAddress%PAGE_SIZE)
  {
    startAddress += PAGE_SIZE - startAddress%PAGE_SIZE;
  }

  Header *holeHeader = (Header *)startAddress;
  holeHeader->size = endAddress-startAddress;
  holeHeader->magic = HEAP_MAGIC;
  holeHeader->isHole = true;
  Footer *holeFooter = (Footer *)(startAddress+holeHeader->size-sizeof(Footer));
  holeFooter->header = holeHeader;
  holeFooter->magic = HEAP_MAGIC;

  // Start by having just one big hole.
  index->insert( (Header *)startAddress );

}


Heap::~Heap()
{
}

s32int Heap::findSmallestHole(u32int size, bool pageAlign)
{
  // Find the smallest hole that will fit.
  u32int iterator = 0;
  while ( iterator < index->getSize())
  {
    // if the user has requested the memory be page-aligned
    if (pageAlign)
    {
      // page-align the starting point of this header.
      u32int location = (u32int)index->getItem(iterator);
      s32int offset = 0;
      if ((location+sizeof(Header))%PAGE_SIZE)
      {
        offset = PAGE_SIZE - (location+sizeof(Header))%PAGE_SIZE;
      }
      
      s32int holeSize = (s32int) (index->getItem(iterator)->size);
      holeSize -= offset;

      // can we fit now?
      if (holeSize >= (s32int)size)
      {
        break;
      }
    }
    else if (index->getItem(iterator)->size >= size)
    {
      break;
    }
  
    iterator++;
  }
  if (iterator == index->getSize())
  {
    return -1;
  }
  else
  {
    return iterator;
  }
}

void *Heap::allocate(u32int size, bool pageAlign)
{
 
  criticalSection();
#ifdef HEAP_DEBUG
  checkIntegrity();
#endif
  
  // Take into account the header/footer.
  u32int newSize = size + sizeof(Header) + sizeof(Footer);

  s32int iterator = findSmallestHole(newSize, pageAlign);

  if (iterator == -1)
  {
    u32int oldLength = endAddress-startAddress;
    u32int oldEndAddress = endAddress;
  
    // A hole big enough was not found. Allocate some more space now.
    expand(endAddress-startAddress+newSize);
    
    u32int newLength = endAddress-startAddress;

    // Find the endmost header. (Not endmost in size, endmost in location)
    u32int iterator2 = 0;
    u32int idx = -1; u32int value = 0x0;
    
    while (iterator2 < index->getSize())
    {
      if ((u32int)index->getItem(iterator2) > value)
      {
        value = (u32int)index->getItem(iterator2);
        idx = iterator2;
      }
      iterator2++;
    }

    // If we didnt find ANY headers, we need to add one.
    if (idx == -1)
    {
      Header *header = (Header *)oldEndAddress;
      header->magic = HEAP_MAGIC;
      header->size = newLength-oldLength;
      header->isHole = true;
      Footer *footer = (Footer *) ((u32int)header + header->size - sizeof(Footer));
      footer->magic = HEAP_MAGIC;
      footer->header = header;
      // Put this new header in the index
      index->insert(header);
    }
    else
    {
      // The last header is the one whose size needs adjusting.
      Header *header = index->getItem(idx);
      header->size += newLength-oldLength;
      // Rewrite it's footer.
      Footer *footer = (Footer *) ( (u32int)header + header->size - sizeof(Footer));
      footer->header = header;
      footer->magic = HEAP_MAGIC;
    }
#ifdef HEAP_DEBUG
    checkIntegrity();
#endif
    endCriticalSection();
    return allocate(size, pageAlign);
  }

  if (index->getItem(iterator)->magic != HEAP_MAGIC)
  {
    kerr << "block: " << hex << (u32int)index->getItem(iterator) << endl;
  }
  ASSERT(index->getItem(iterator)->magic == HEAP_MAGIC);
  
  u32int origHolePos = (u32int)index->getItem(iterator);
  u32int origHoleSize= (u32int)index->getItem(iterator)->size;

  // If the original hole size - the requested hole size is less than
  // the space required to make a new hole (sizeof(header)+sizeof(footeR)),
  // then just use the origHoleSize.
  if (origHoleSize-newSize < sizeof(Header)+sizeof(Footer))
  {
    size += (origHoleSize-newSize);
    newSize = origHoleSize;
  }
  
  // If we need to page-align the data, do it now and make a new hole.
  if (pageAlign && origHolePos%PAGE_SIZE)
  {
    u32int newLocation = origHolePos + PAGE_SIZE - origHolePos%PAGE_SIZE - sizeof(Header);
    Header *holeHeader = (Header *)origHolePos;
    holeHeader->size   = PAGE_SIZE - origHolePos%PAGE_SIZE - sizeof(Header);
    holeHeader->magic  = HEAP_MAGIC;
    holeHeader->isHole = true;
    Footer *holeFooter = (Footer *) ((u32int)newLocation-sizeof(Footer));
    holeFooter->magic  = HEAP_MAGIC;
    holeFooter->header = holeHeader;
    origHolePos        = newLocation;
    origHoleSize       = origHoleSize - holeHeader->size;
  }
  else
  {
    // Delete the hole.
    index->remove(iterator);
  }
//   kerr << "Assigning: " << hex << newSize-sizeof(Header)-sizeof(Footer) << ", " << newSize << ", " << origHolePos << endl;
  // Overwrite the original header.
  Header *blockHeader  = (Header *)origHolePos;
  blockHeader->magic   = HEAP_MAGIC;
  blockHeader->isHole  = false;
  blockHeader->size    = newSize;
  if (processManager->getProcess())
  {
    blockHeader->pid   = processManager->getProcess()->getPid();
    for(int i = 0; i < NBACKTRACE; i++)
    {
      blockHeader->backtrace[i] = Kernel::backtrace(i);
//       kerr << hex << blockHeader->backtrace[i] << ", ";
    }
  }
  else
  {
    blockHeader->pid   = 0;
    for(int i = 0; i < NBACKTRACE; i++)
    {
      blockHeader->backtrace[i] = 0;
    }
  }
//   kerr << endl;
  // And the footer...
  Footer *blockFooter  = (Footer *) (origHolePos+sizeof(Header)+size);
  blockFooter->magic   = HEAP_MAGIC;
  blockFooter->header  = blockHeader;
  
  // If the new hole wouldn't have size zero...
  if (origHoleSize-newSize)
  {
    // Write it.
    Header *holeHeader = (Header *) (origHolePos+sizeof(Header)+size+sizeof(Footer));
    holeHeader->magic  = HEAP_MAGIC;
    holeHeader->isHole = true;
    holeHeader->size   = origHoleSize-newSize;
    
    Footer *holeFooter = (Footer *) ( (u32int)holeHeader+
                                     origHoleSize-newSize-sizeof(Footer) );
    if ((Address)holeFooter < 0xC0000000 || (Address)holeFooter > 0xE0000000)
    {
      kerr << "Footer: " << hex << (u32int)holeFooter << endl;
      kerr << "origHoleSize: " << hex << origHoleSize << endl;
      kerr << "newSize: " << hex << newSize << endl;
      kerr << "header: " << hex << (u32int)holeHeader << endl;
    }
    if ((Address)holeFooter < endAddress)
    {
      holeFooter->magic  = HEAP_MAGIC;
      holeFooter->header = holeHeader;
    }
    
    // Put the new hole in the index.
    index->insert(holeHeader);
  }
  
#ifdef HEAP_DEBUG
  checkIntegrity();
#endif
  endCriticalSection();
  
  return (void *) ( (u32int)blockHeader+sizeof(Header) );
  
}

void Heap::free(void *p)
{
  criticalSection();
#ifdef HEAP_DEBUG
  checkIntegrity();
#endif
  
  // Exit gracefully for null pointers.
  if (!p)
  {
    return;
  }

  // Get the header associated with this pointer.
  Header *header = (Header *) ( (u32int)p - sizeof(Header) );
  Footer *footer = (Footer *) ( (u32int)header + header->size - sizeof(Footer) );
  
  // Consistency check...
  if (header->magic != HEAP_MAGIC)
  {
    kerr << "Header: " << (u32int)header << ", magic: " << header->magic << endl;
    kerr << "size: " << header->size <<endl;
  }
  ASSERT(header->magic == HEAP_MAGIC);
  ASSERT(footer->magic == HEAP_MAGIC);
  ASSERT(!header->isHole);
  
  // Holeify us
  header->isHole = true;
  
  // Do we want to add the header into the index?
  bool doAdd = true;
  
  // Unify left
  // if the thing immediately to the left of us is a footer...
  Footer *testFooter = (Footer *) ( (u32int)header - sizeof(Footer) );
  if (testFooter->magic == HEAP_MAGIC &&
      testFooter->header->isHole )
  {
    // cache our current size.
    u32int cacheSize = header->size;
    
    // rewrite our header with the new one
    header = testFooter->header;
    
    // rewrite our footer to point to the new header.
    footer->header = header;
    
    // change the size.
    header->size += cacheSize;
    
    // Since this header is already in the index, we don't want to add it again.
    doAdd = false;
  }
  

  // Unify right
  // if the the thing immediately to the right of us is a header...
  Header *testHeader = (Header *) ( (u32int)footer + sizeof(Footer) );
  if (testHeader->magic == HEAP_MAGIC &&
      testHeader->isHole )
  {
    // increase our size.
    header->size += testHeader->size;
    
    // rewrite it's footer to point to our header.
    testFooter = (Footer *) ( (u32int)testHeader + testHeader->size
                                      - sizeof(Footer));
    testFooter->header = header;
    
    // It's now OUR footer! muahahaha....
    footer = testFooter;
    
    // find and remove this header from the index.
    u32int iterator = 0;
    while ( (iterator < index->getSize()) &&
            (index->getItem(iterator) != testHeader) )
    {
      iterator ++;
    }

    // Make sure we actually found the item.
    ASSERT(iterator < index->getSize());
    
    // Remove it.
    index->remove(iterator);
  }

  // If the footer location is the end address, we can contract.
  if ( (u32int)footer+sizeof(Footer) == endAddress )
  {
    u32int oldLength = endAddress-startAddress;
    u32int newLength = contract((u32int)header - startAddress);
    
    // Check how big we will be after resizing
    if (header->size - (oldLength-newLength)) // if we are > 0
    {
      // we still exist, resize us.
      header->size-= oldLength-newLength;
      footer = (Footer *) ( (u32int)header + header->size - sizeof(Footer) );
      footer->magic = HEAP_MAGIC;
      footer->header = header;
    }
    else
    {
      // We no longer exist :(. Remove us from the index.
      u32int iterator = 0;
      while ( (iterator < index->getSize()) &&
              (index->getItem(iterator) != header) )
      {
        iterator ++;
      }
      
      // If we didnt find ourselves, we have nothing to remove.
      if (iterator < index->getSize())
      {
        index->remove(iterator);
      }
    }
#ifdef HEAP_DEBUG
    checkIntegrity();
#endif
  }
  // Add us to the index
  if (doAdd)
  {
    index->insert(header);
  }
  
#ifdef HEAP_DEBUG
  checkIntegrity();
#endif
  endCriticalSection();

}

void Heap::expand(u32int newSize)
{
#ifdef HEAP_DEBUG
  checkIntegrity();
#endif
  // Sanity check.
  ASSERT(newSize > endAddress-startAddress);
  
  DEBUG_MSG("Heap expanding from " << hex << endAddress-startAddress << " to " << newSize);
  
  // Get the nearest following page boundary.
  if (newSize%PAGE_SIZE)
  {
    newSize &= PAGE_MASK;
    newSize += PAGE_SIZE;
  }
  
  // Make sure we are not overreaching ourselves.
  ASSERT(startAddress+newSize <= maxAddress);
  
  // This should always be on a page boundary.
  u32int oldSize = endAddress-startAddress; 
  
  u32int i = oldSize;
  while(i < newSize)
  {
    memoryManager.allocFrame( memoryManager.kernelDirectory->
                              getPage(startAddress+i), isKernel );
    i += PAGE_SIZE;
  }
  
  endAddress = startAddress+newSize;
#ifdef HEAP_DEBUG
//   checkIntegrity();
#endif
}

u32int Heap::contract(u32int newSize)
{
#ifdef HEAP_DEBUG
  checkIntegrity();
#endif
  // Sanity check.
  ASSERT(newSize < endAddress-startAddress);
  
  // get the nearest following page boundary.
  if (newSize%PAGE_SIZE)
  {
    newSize += PAGE_SIZE - newSize%PAGE_SIZE;
  }
  
  // Don't contract too far.
  if (newSize < HEAP_MIN_SIZE)
    newSize = HEAP_MIN_SIZE;
  
  DEBUG_MSG("Heap contracting from " << hex << endAddress-startAddress << " to " << newSize);
  
  // Make sure we are not overreaching ourselves.
  ASSERT(newSize > 0);
  
  u32int oldSize = endAddress-startAddress;
  
  u32int i = newSize;
  while(i < oldSize)
  {
    memoryManager.freeFrame( memoryManager.kernelDirectory->
                             getPage(startAddress+i) );
    i += PAGE_SIZE;
  }
  
  endAddress = startAddress+newSize;
#ifdef HEAP_DEBUG
//   checkIntegrity();
#endif
  return newSize;
}

void Heap::checkIntegrity()
{
  // We should, by starting at startAddress, be able to walk through all blocks/
  // holes and check their magic numbers.
  u32int addr = startAddress;
  Header *lastHeader = NULL;
  Header *thisHeader = (Header*)startAddress;
  Header *nextHeader = (Header*)( (u32int)thisHeader+thisHeader->size );
  if ((u32int)nextHeader >= endAddress)
    nextHeader = NULL;
  while (thisHeader)
  {
    if (thisHeader->magic != HEAP_MAGIC)
    {
      // header overwritten.
      kerr << "\nPrevious block\t\tBacktrace\tBacktrace(cont)" << endl;
      kerr << "Address: " << hex << (u32int)lastHeader << "\t"
           << lastHeader->backtrace[0] << "\t" << lastHeader->backtrace[3] << endl;
      kerr << "Size: " << lastHeader->size << "\t\t" << lastHeader->backtrace[1]
           << "\t" << lastHeader->backtrace[4] << endl;
      kerr << "Pid: " << lastHeader->pid << "\t\t" << lastHeader->backtrace[2]
           << "\t" << lastHeader->backtrace[5] << endl;
      kerr << "Hole: " << ((lastHeader->isHole)?1:0) << endl;
      kerr << "\nThis block\t\tBacktrace\tBacktrace(cont)" << endl;
      kerr << "Address: " << hex << (u32int)thisHeader << "\t"
          << thisHeader->backtrace[0] << "\t" << thisHeader->backtrace[3] << endl;
      kerr << "Size: " << thisHeader->size << "\t\t" << thisHeader->backtrace[1]
          << "\t" << thisHeader->backtrace[4] << endl;
      kerr << "Pid: " << thisHeader->pid << "\t\t" << thisHeader->backtrace[2]
          << "\t" << thisHeader->backtrace[5] << endl;
      kerr << "Hole: " << ((thisHeader->isHole)?1:0) << endl;
      PANIC("Heap header overwritten!");
    }
    
    if (!nextHeader)
      break;
    Footer *footer = (Footer*)((u32int)nextHeader-sizeof(Footer));

    if (footer->magic != HEAP_MAGIC)
    {
      // footer overwritten.
      kerr << "\nPrevious\tThis\t\tNext" << endl;
      kerr << hex << (u32int)lastHeader << "\t" << (u32int)thisHeader
          << "\t" << (u32int)nextHeader << endl;
      kerr << hex << lastHeader->size << "\t\t" << thisHeader->size << "\t\t"
          << nextHeader->size << endl;
      kerr << hex << lastHeader->pid << "\t\t" << thisHeader->pid << "\t\t"
          << nextHeader->pid << endl;
      PANIC("Heap footer overwritten!");
    }
    
    lastHeader = thisHeader;
    thisHeader = nextHeader;
    nextHeader = (Header*)( (u32int)thisHeader+thisHeader->size );
    if ((u32int)nextHeader >= endAddress)
      nextHeader = NULL;
  }
  
}
